package com.autoagent.ai_live_master.common.utils;

import com.autoagent.ai_live_master.common.config.ApiConfig;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Data
@AllArgsConstructor
@NoArgsConstructor
public class KbClient {
    private static final String UPLOAD_ENDPOINT = "/openapi/fs/upload";
    private static final String CREATE_DATASET_ENDPOINT = "/openapi/kb/createDsAndTask";
    private static final String MODIFY_DATASET_ENDPOINT = "/openapi/kb/modifyDsAndTask";
    private static final String DELETE_DATASET_ENDPOINT = "/openapi/kb/deleteDs";
    private static final String QUERY_DATASET_ENDPOINT = "/api/kb/ds/query";
    private static final String AUTH_ENDPOINT = "/openapi/user/auth";


    private ApiConfig apiConfig;
    private String authKey;
    private String authSecret;


    /**
     * 上传文件到知识库
     * @param fileContent 文件内容字节数组
     * @param fileName 文件名
     * @param returnType 返回类型(id/url)
     * @param metadata 文件元数据
     * @return 文件ID或URL
     * @throws Exception 上传失败时抛出异常
     */
    public String uploadFile(byte[] fileContent, String fileName, String returnType, Map<String, Object> metadata) throws Exception {
        String boundary = "----WebKitFormBoundary" + System.currentTimeMillis();
        String url = apiConfig.getHost() + UPLOAD_ENDPOINT;
        
        // 构建multipart请求体
        ByteArrayOutputStream requestBody = new ByteArrayOutputStream();
        
        // 添加文件部分
        requestBody.write(("--" + boundary + "\r\n").getBytes(StandardCharsets.UTF_8));
        requestBody.write(("Content-Disposition: form-data; name=\"file\"; filename=\"" + fileName + "\"\r\n").getBytes(StandardCharsets.UTF_8));
        requestBody.write("Content-Type: application/octet-stream\r\n\r\n".getBytes(StandardCharsets.UTF_8));
        requestBody.write(fileContent);
        requestBody.write("\r\n".getBytes(StandardCharsets.UTF_8));
        
        // 添加returnType参数
        requestBody.write(("--" + boundary + "\r\n").getBytes(StandardCharsets.UTF_8));
        requestBody.write("Content-Disposition: form-data; name=\"returnType\"\r\n\r\n".getBytes(StandardCharsets.UTF_8));
        requestBody.write((returnType + "\r\n").getBytes(StandardCharsets.UTF_8));
        
        // 添加metadata参数(如果存在)
        if (metadata != null) {
            requestBody.write(("--" + boundary + "\r\n").getBytes(StandardCharsets.UTF_8));
            requestBody.write("Content-Disposition: form-data; name=\"metadata\"\r\n\r\n".getBytes(StandardCharsets.UTF_8));
            requestBody.write((new ObjectMapper().writeValueAsString(metadata) + "\r\n").getBytes(StandardCharsets.UTF_8));
        }
        
        // 结束边界
        requestBody.write(("--" + boundary + "--\r\n").getBytes(StandardCharsets.UTF_8));
        
        // 发送请求
        HttpClient client = getHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Authorization", String.format("Bearer %s.%s", this.authKey, this.authSecret))
                .header("Content-Type", "multipart/form-data; boundary=" + boundary)
                .POST(HttpRequest.BodyPublishers.ofByteArray(requestBody.toByteArray()))
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

        // 处理响应
        if (response.statusCode() != 200) {
            throw new Exception("上传失败: " + response.body());
        }

        JsonNode result = new ObjectMapper().readTree(response.body());
        if (result.get("code").asInt() != 1) {
            throw new Exception(result.get("msg").asText());
        }
        
        return result.get("data").asText();
    }

    /**
     * 创建数据集并开始任务(支持可选参数)
     * @param kbId 知识库ID(必需)
     * @param name 数据集名称(必需)
     * @param fileId 文件ID(必需)
     * @param options 可选参数配置
     * @return 数据集ID
     * @throws Exception 如果创建过程中发生错误
     */
    public Long createDataset(int kbId, String name, String fileId, DatasetOptions options) throws Exception {
        // 构建请求数据
        Map<String, Object> data = new HashMap<>();
        data.put("kbId", kbId);
        data.put("fileId", fileId);
        data.put("name", name);
        data.put("parserType", options != null ? options.getParserType() : "general");
        data.put("skipIfExists", options == null || options.isSkipIfExists());
        data.put("autoStartChunkTask", options == null || options.isAutoStartChunkTask());
        data.put("alwaysCreateNew", options != null && options.isAlwaysCreateNew());

        // 直接构造HTTP请求
        String url = apiConfig.getHost() + CREATE_DATASET_ENDPOINT;
        ObjectMapper mapper = new ObjectMapper();
        String jsonPayload = mapper.writeValueAsString(data);
        
        HttpClient client = getHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Authorization", String.format("Bearer %s.%s", this.authKey, this.authSecret))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(jsonPayload))
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

        // 处理响应
        if (response.statusCode() != 200) {
            throw new Exception("创建数据集失败: " + response.body());
        }

        JsonNode result = mapper.readTree(response.body());
        if (result.get("code").asInt() != 1) {
            throw new Exception(result.get("msg").asText());
        }
        
        return result.get("data").asLong();
    }

    // 可选参数配置类
    @Data
    @Builder
    public static class DatasetOptions {
        @Builder.Default
        private String parserType = "general";
        @Builder.Default
        private boolean skipIfExists = true;
        @Builder.Default
        private boolean autoStartChunkTask = true;
        @Builder.Default
        private boolean alwaysCreateNew = false;
    }



    /**
     * 上传文件到知识库并创建数据集
     * @param kbId 知识库ID
     * @param fileContent 文件内容字节数组
     * @param fileName 文件名
     * @return 包含文件ID和数据集ID的Map
     * @throws Exception 如果操作过程中发生错误
     */
    public Map<String, Object> uploadFileToKb(int kbId, byte[] fileContent, String fileName) throws Exception {
        try {
            // 上传文件
            String fileId = uploadFile(fileContent, fileName, "id", null);
            
            // 创建数据集
            Long datasetId = createDataset(kbId, fileName, fileId, null);
            
            // 返回结果
            Map<String, Object> result = new HashMap<>();
            result.put("code", 200);
            
            Map<String, Object> data = new HashMap<>();
            data.put("fileId", fileId);
            data.put("datasetId", datasetId);
            result.put("data", data);
            
            return result;
        } catch (Exception e) {
            throw new Exception("上传文件到知识库失败: " + e.getMessage());
        }
    }

    /**
     * 删除数据集
     * @param dsIds 数据集ID列表
     * @return 是否删除成功
     * @throws Exception 如果删除过程中发生错误
     */
    public boolean deleteDataset(List<Integer> dsIds) throws Exception {
        // 构建请求数据
        Map<String, Object> data = new HashMap<>();
        data.put("dsIdList", dsIds);

        // 发送请求
        String url = apiConfig.getHost() + DELETE_DATASET_ENDPOINT;
        ObjectMapper mapper = new ObjectMapper();
        String jsonPayload = mapper.writeValueAsString(data);
        
        HttpClient client = getHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Authorization", String.format("Bearer %s.%s", this.authKey, this.authSecret))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(jsonPayload))
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

        // 处理响应
        if (response.statusCode() != 200) {
            throw new Exception("删除数据集失败: " + response.body());
        }

        JsonNode result = mapper.readTree(response.body());
        if (result.get("code").asInt() != 1) {
            throw new Exception(result.get("msg").asText());
        }
        
        return result.get("data").asBoolean();
    }
    
    /**
     * 查询知识库中的数据集
     * @param kbId 知识库ID
     * @param pageNum 页码
     * @param pageSize 每页数量
     * @param count 是否统计总数
     * @param keywords 搜索关键词
     * @param parentId 父级ID
     * @return 查询结果的 JsonNode（包含分页信息和数据列表）
     * @throws Exception 请求失败时抛出异常
     */
    public JsonNode queryDatasets(int kbId, int pageNum, int pageSize, Boolean count, String keywords, Integer parentId) throws Exception {
        String token = fetchJwtToken(); // 先获取token

        String url = apiConfig.getHost() + QUERY_DATASET_ENDPOINT;
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("pageNum", pageNum);
        requestBody.put("pageSize", pageSize);
        requestBody.put("kbId", kbId);
        requestBody.put("count", count != null ? count : true);

        if (keywords != null && !keywords.trim().isEmpty()) {
            requestBody.put("keywords", keywords);
        }

        if (parentId != null) {
            requestBody.put("parentId", parentId);
        }

        String jsonPayload = new ObjectMapper().writeValueAsString(requestBody);

        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Authorization", "Bearer " + token)  // 用获取到的token
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(jsonPayload))
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

        if (response.statusCode() != 200) {
            throw new Exception("查询数据集失败: " + response.body());
        }

        JsonNode result = new ObjectMapper().readTree(response.body());
        if (result.get("code").asInt() != 1) {
            throw new Exception("查询数据集失败: " + result.get("msg").asText());
        }

        return result.get("data");
    }

    private HttpClient getHttpClient() {
        return HttpClient.newHttpClient();
    }

    public String fetchJwtToken() throws Exception {
        String url = apiConfig.getHost() + AUTH_ENDPOINT;
        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Authorization", String.format("Bearer %s.%s", this.authKey, this.authSecret))
                .GET()
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());

        if (response.statusCode() != 200) {
            throw new Exception("获取JWT Token失败: " + response.body());
        }

        JsonNode json = new ObjectMapper().readTree(response.body());
        if (json.get("code").asInt() != 1) {
            throw new Exception("获取JWT Token失败: " + json.get("msg").asText());
        }

        return json.get("data").get("token").asText();
    }

}
